import argparse
import os
from utils.angr_util import AngrUtil
from tqdm import tqdm
import pickle


def get_limited_data(file_list, max=100):
    tmp = []
    count_dict = {}
    for f in file_list:
        split_names = f.split("-")
        if len(split_names) < 2:
            continue
        is_auto = split_names[0] == "auto_schedule"
        if not is_auto:
            continue
        func_name = split_names[0]
        if func_name not in count_dict:
            count_dict[func_name] = 0
        if count_dict[func_name] >= max:
            continue
        tmp.append(f)
        count_dict[func_name] += 1
    return tmp


def run(args):
    path = args.path
    if not os.path.exists(path):
        raise Exception("Dir Not Found!")

    if not os.path.isdir(path):
        raise Exception("Not A Dir!")

    file_list = os.listdir(path)

    if args.limit > -1:
        file_list = get_limited_data(file_list, args.limit)
    is_auto = False

    trace_data = {}
    if args.trace_data != "None":
        f = open(args.trace_data, "rb")
        trace_data = pickle.load(f)
        f.close()

    ignore_list = []
    f = open(args.ignore_config, "r")
    for line in f.readlines():
        i = int(line.split(":")[0])
        ignore_list.append(i)
    f.close()

    pbar = tqdm(total=len(file_list))

    for idx, filename in enumerate(file_list):
        pbar.update(1)

        split_names = filename.split("-")
        if len(split_names) < 2:
            continue
        is_auto = split_names[0] == "auto_schedule"
        if not is_auto:
            continue

        if idx in ignore_list:
            continue

        if filename in trace_data:
            continue

        print("\n${}?{}\n".format(idx, filename))

        util = AngrUtil("{}/{}".format(path, filename))
        util.identify_function()
        trace_data[filename] = {}
        for function_name in util.function_names:
            trace_data[filename][function_name] = {}
            trace_data[filename][function_name]['trace'] = []
            trace_data[filename][function_name]['blocks'] = []

            g = util.get_function_graph(function_name)
            blocks = util.get_func_blocks_by_name(function_name)
            blocks_dict = util.contruct_block_dict(blocks)
            if args.extend == 1:
                g, blocks_dict = util.try_extend_graph(g, blocks_dict)

            for k in blocks_dict:
                tokens = util.get_block_tokens(blocks_dict[k])
                trace_data[filename][function_name]['blocks'].append(tokens)

            for _ in range(args.trace_number):
                trace = util.random_walk(blocks_dict, g)
                tokens = util.parsing_trace(trace)
                # trace_data.append(tokens)
                trace_data[filename][function_name]['trace'].append(tokens)

        f = open("as_trace_data_all.pkl", "wb")
        if args.trace_data != "None":
            f.close()
            f = open(args.trace_data, "wb")
        pickle.dump(trace_data, f)
        f.close()

    pbar.close()

    # print("Generating the Word2Vec Model ...")

    # token_list = []
    # for file in trace_data:
    #     funcs = trace_data[file]
    #     for func in funcs:
    #         tmp_trace = trace_data[file][func]['trace']
    #         token_list.extend(tmp_trace)

    # model = Word2Vec(sentences=token_list, vector_size=200,
    #                  window=10, min_count=1, workers=4)
    # model.save(args.output)
    # print("Finish!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Generate the trace from CFG")
    parser.add_argument("-p", "--path",   type=str,
                        help="The path of .so data")
    parser.add_argument("-e", "--extend", type=int, default=0,
                        help="The trigger of extending the graph")
    parser.add_argument("-t", "--trace_number", type=int,
                        default=10, help="The number of traces")
    parser.add_argument("-o", "--output", type=str,
                        default="word2vec-train.model",
                        help="The name of the model")
    parser.add_argument("-i", "--ignore_config", type=str,
                        default="ignore_function.txt",
                        help="The config for ignoring files")
    parser.add_argument("--trace_data", type=str,
                        default="None", help="The path of trace data")
    parser.add_argument("-l", "--limit", type=int,
                        default=-1, help="The maximum of function files")

    args = parser.parse_args()
    run(args)
